# coding: utf-8
#### 所有的层
from common.functions import *
from common.util import *

# 普通层的激活函数
class Relu:
    def __init__(self):
        self.mask = None

    def forward(self, x):
        self.mask = (x <= 0)  # 标记元素值不大于0的索引
        out = x.copy()
        out[self.mask] = 0

        return out

    def backward(self, dout):
        dout[self.mask] = 0
        dx = dout

        return dx

# 普通层的激活函数（另一个选择）
class Sigmoid:
    def __init__(self):
        self.out = None

    def forward(self, x):
        out = sigmoid(x)
        self.out = out
        return out

    def backward(self, dout):
        dx = dout * (1 - self.out) * self.out  # 觉得这里奇怪，就去看书上的推导，很清晰

        return dx


# 矩阵计算层
class Affine:
    def __init__(self, W, b):
        self.W = W
        self.b = b

        self.x = None  # 哦哦，原来把这个输入存下来，是为了在计算dw的时候用啊
        self.orginal_x_shape = None  # 张量？

        self.dW = None
        self.db = None

    def forward(self, x):
        # 对应张量
        self.orginal_x_shape = x.shape
        x = x.reshape(x.shape[0], -1)  # 奇奇怪怪？
        self.x = x

        out = np.dot(self.x, self.W) + self.b

        return out

    def backward(self, dout):
        dx = np.dot(dout, self.W.T)
        self.dW = np.dot(self.x.T, dout)  # 反向的时候不是主要求这两个值吗？
        self.db = np.sum(dout, axis=0)

        # ???? 我上面的dot操作不是已经还原形状了吗？
        dx = dx.reshape(*self.orginal_x_shape)  # 还原输入数据的形状（对应张量）
        return dx


# 输出层的激活函数连带损失函数一起，把这两层组起来之后，反向的值很好看
class SoftmaxWithLoss:
    def __init__(self):
        self.loss = None
        self.y = None  # softmax的输出
        self.t = None  # 监督数据

    def forward(self, x, t):
        self.t = t
        self.y = softmax(x)
        self.loss = cross_entropy_error(self.y, self.t)

        return self.loss  # 对的，学习的最终目的就是求解损失函数的值

    def backward(self, dout=1):
        batch_size = self.t.shape[0]  # 学习数据的数量
        if self.t.size == self.y.size:  # 监督数据是one-hot-vector的情况
            dx = (self.y - self.t) / batch_size
        else:
            dx = self.y.copy()
            dx[np.arange(batch_size), self.t] -= 1  # 懂了，设想一下这样处理之后，和上面的处理是一样的结果，为什么只有正确标签对应的那个元素减1呢？想一下完整的t标签里，只有一个1，其他全部都是0，那么减去0，和不做任何改变，有区别吗？
            dx = dx / batch_size

        return dx

class Convolution:
    def __init__(self, W, b, stride=1, pad=0):
        self.W = W
        self.b = b
        self.stride = stride
        self.pad = pad

        # 中间数据（backward时使用？）
        self.x = None
        self.col = None
        self.col_W = None

        # 权重和偏置参数的梯度
        self.dW = None
        self.db = None

    def forward(self, x):
        FN, C, FH, FW = self.W.shape  # 获取滤波器（权重）的数量，通道数，滤波器的高，滤波器的宽
        N, C, H, W = x.shape  # 获取输入数据的 数量、通道数、高、宽
        out_h = 1 + int((H + 2*self.pad - FH) / self.stride)  # 通过公式先提前计算出，输出数据的形状
        out_w = 1 + int((W + 2*self.pad - FW) / self.stride)  # 同上

        col = im2col(x, FH, FW, self.stride, self.pad)  # 通过im2col将四维数组转换为二维的
        col_W = self.W.reshape(FN, -1).T  # 滤波器同样降为二维数组

        out = np.dot(col, col_W) + self.b  # 像往常一样计算权重和
        out = out.reshape(N, out_h, out_w, -1).transpose(0, 3, 1, 2)  # 前面重新定义形状还好说，后面换轴干嘛？
        # 哦哦，针对上一句的大概性理解应该是这样的，我计算到的out肯定是对的数据，但是它是二维的，维度不对，我需要调一下
        # 我本来应该把它调成：数量、通道数、高、宽这样的形式，对的，我理应将它调成这样的形式，但是，我所能知道的就只有
        # 数量、高、宽三个维度数据，所以，我需要借助一下reshape这个函数，通过设定好reshape(数量、高、宽、-1）这样的形式
        # 我可以求出通道数的维度应该填多少，那么此时，out的形状就为：数量、高、宽、通道数，然后我们再利用transpose这个函数
        # 把out的形状调整为标准的，out：数量、通道数、高、宽，对，差不多就是这样

        # 存好这三个数，后面backward应该用的上
        self.x = x
        self.col = col
        self.col_W = col_W

        return out
    # backward 后面再说
    def backward(self, dout):
        FN, C, FH, FW = self.W.shape
        dout = dout.transpose(0,2,3,1).reshape(-1, FN)

        self.db = np.sum(dout, axis=0)
        self.dW = np.dot(self.col.T, dout)
        self.dW = self.dW.transpose(1, 0).reshape(FN, C, FH, FW)

        dcol = np.dot(dout, self.col_W.T)
        dx = col2im(dcol, self.x.shape, FH, FW, self.stride, self.pad)

        return dx


# 池化层
class Pooling:
    def __init__(self, pool_h, pool_w, stride=1, pad=0):
        self.pool_h = pool_h
        self.pool_w = pool_w
        self.stride = stride
        self.pad = pad

        self.x = None
        self.arg_max = None

    def forward(self, x):
        N, C, H, W = x.shape
        out_h = int(1 + (H - self.pool_h) / self.stride)  # 咦，这里不加填充了是吧
        out_w = int(1 + (W - self.pool_w) / self.stride)  # 咦，这里不加填充了是吧

        col = im2col(x, self.pool_h, self.pool_w, self.stride, self.pad)
        col = col.reshape(-1, self.pool_h*self.pool_w)

        arg_max = np.argmax(col, axis=1)  # 第一维方向上取最大值
        out = np.max(col, axis=1)  # ?和上面重复操作？，所以是，上面那个是后面要用到，所以单独放了？
        out = out.reshape(N, out_h, out_w, C).transpose(0, 3, 1, 2)  # 离谱，一定有某种原因让它的操作为先重置形状然后再换轴，可恶，但是我想不通

        self.x = x
        self.arg_max = arg_max

        return out

    # backward后面再说！
    def backward(self, dout):
        dout = dout.transpose(0, 2, 3, 1)

        pool_size = self.pool_h * self.pool_w
        dmax = np.zeros((dout.size, pool_size))
        dmax[np.arange(self.arg_max.size), self.arg_max.flatten()] = dout.flatten()
        dmax = dmax.reshape(dout.shape + (pool_size,))

        dcol = dmax.reshape(dmax.shape[0] * dmax.shape[1] * dmax.shape[2], -1)
        dx = col2im(dcol, self.x.shape, self.pool_h, self.pool_w, self.stride, self.pad)

        return dx









